"""
model name : 深度学习
file       : YOLO_LRScheduler.py
information:
    author : OuYang
    time   : 2025/1/16
"""


class YOLOLRScheduler:
    def __init__(self, optimizer, last_epoch=-1):
        self.optimizer = optimizer
        self.last_epoch = last_epoch
        self.base_lrs = [1e-3 for _ in range(75)] + \
                        [1e-3 for _ in range(30)] + \
                        [1e-4 for _ in range(30)]

        # 1       1e-3
        # 2-75    1e-2
        # 76-105  1e-3
        # 106-135 1e-4

    def get_lr(self):
        self.last_epoch += 1
        return self.base_lrs[self.last_epoch]

    def step(self):
        for param_group in self.optimizer.param_groups:
            param_group['lr'] = self.get_lr()


print([1e-3 for _ in range(75)] + \
      [1e-4 for _ in range(30)] + \
      [1e-5 for _ in range(30)])
